#ifndef _HEAD_BUFFER_H_
#define _HEAD_BUFFER_H_

#include <cmath>
#include <vector>
#include <iostream>
#include <eigen3/Eigen/Dense>

#include <std_msgs/Float64MultiArray.h>

class HeadBuffer
{
public:
  // Resolution: 22.5 degrees one piece. MID=7.
  // Yaw=0. Larger: 7 - 15 :  0 to Pi;
  // Smaller : 7 - 0 : 0 to(nearly) - Pi;
  static constexpr int LOOKING_PIECES_SIZE = 16;

  struct HeadPlanParams
  {
    double k_current_v = 1.0;
    double k_planned_direction = 0.4;
    double k_v_fluctuation = 0.6;
    double k_dynamic_objects = 1.0;
    double k_common_update = 0.15;
  } hp;

  HeadBuffer() {}
  ~HeadBuffer() {}

  bool setParams(
      double fov,
      double max_v_xy,
      HeadPlanParams &_hp)
  {
    MAX_V_XY = max_v_xy;
    hp = _hp;

    return setFov(fov);
  }

  bool setFov(double fov)
  {
    int valid_piece_num = (int)((double)fov / 180.0 * M_PI / heading_resolution);
    if (valid_piece_num % 2 == 0) valid_piece_num -= 1;
    if (valid_piece_num < 1) {
      std::cerr << "No enough view field with the current camera set!" << std::endl;
      return false;
    }
    valid_piece_num_one_side = (valid_piece_num - 1) / 2;

    std::cout << "Heading resolution = " << heading_resolution
              << " (rad), Fov = " << fov
              << ", valid_piece_num_one_side = " << valid_piece_num_one_side << std::endl;

    return true;
  }

  void setVelocity(const Eigen::Vector3d &v) { v0 = v; }
  void setDynamicObjectsYaw(const std::vector<double> &yaw)
  {
    dynamic_objects_yaw.clear();
    dynamic_objects_yaw = yaw;
  }

  void updateHeadBuffer(double yaw, const Eigen::Vector3d &v)
  {
    setVelocity(v);

    int heading_direction_seq = getHeadingSeq(yaw);

    Eigen::VectorXi update_flag_buf = Eigen::VectorXi::Zero(LOOKING_PIECES_SIZE);

    // Add hit for the direction in FOV
    addHitOnePiece(heading_direction_seq);
    update_flag_buf(heading_direction_seq) = 1;

    for (int i = 1; i <= valid_piece_num_one_side; i++) {
      int seq_this_1 = heading_direction_seq + i;
      int seq_this_2 = heading_direction_seq - i;

      correctPieceSeq(seq_this_1);
      correctPieceSeq(seq_this_2);

      addHitOnePiece(seq_this_1);
      addHitOnePiece(seq_this_2);

      update_flag_buf(seq_this_1) = 1;
      update_flag_buf(seq_this_2) = 1;
    }
    // Add miss for the rest
    for (int j = 0; j < LOOKING_PIECES_SIZE; j++) {
      if (update_flag_buf(j) != 1) {
        addMissOnePiece(j);
      }
    }
  }

  /**
   * @brief Get the Optimal Yaw object
   *
   * @param collision_check_pass_flag  >0: pass, <=0: not pass
   * @param theta_h_chosen :the theta angle of 1.5m on the planned path
   * @param if_publish_panel_arrays
   * @return double
   */
  double getOptimalYaw(int collision_check_pass_flag,
                       double theta_h_chosen,
                       bool if_publish_panel_arrays = false)
  {
    // Current velocity direction(v_direction), planned velocity direction(theta_h_chosen),
    // yaw of the head must be in the same coordinate!!
    double yaw_to_send;

    /// Larger update degree indicates the influence of this item
    double k_current_v = 0;

    if (fabs(v0(0)) > 0.05 || fabs(v0(1) > 0.05)) { // if velocity is large, use this cost from velocity direction
      // Give a discount on hp.k_current_v (max) so that when the velocity of the drone is low, the influence of the velocity direction is small.
      k_current_v = std::max(std::min(std::max(fabs(v0(0) / MAX_V_XY), fabs(v0(1) / MAX_V_XY)), 1.0), 0.0);
      v_direction = atan2(v0(1), v0(0));
    }
    double coefficient_current_v = k_current_v * (1.0 - _direction_update_buffer(getHeadingSeq(v_direction)));

    double coefficient_planned_dir = 0.0;
    if (collision_check_pass_flag > 0) coefficient_planned_dir = 1.0 - _direction_update_buffer(getHeadingSeq(theta_h_chosen));

    double min_head_plan_cost = std::numeric_limits<double>::max(); // smaller cost is better
    for (int i = 0; i < LOOKING_PIECES_SIZE; i++) {
      double head_yaw_plan_temp = getHeadingYawFromSeq(i);

      double cost_current_velocity = hp.k_current_v * coefficient_current_v * (v_direction - head_yaw_plan_temp) * (v_direction - head_yaw_plan_temp);

      double cost_update_degree = hp.k_common_update * _direction_update_buffer(i);

      double cost_planned_direction = hp.k_planned_direction * coefficient_planned_dir * (theta_h_chosen - head_yaw_plan_temp) * (theta_h_chosen - head_yaw_plan_temp);
      double cost_head_fluctuation = hp.k_v_fluctuation * (head_yaw_plan_temp - last_head_yaw_plan) * (head_yaw_plan_temp - last_head_yaw_plan);

      double cost_dynamic_objects = 0.0;
      double cost_dynamic_objects_min = -hp.k_dynamic_objects * 2;

      static int effect_range_one_side = 3;
      static double cost_one_piece = hp.k_dynamic_objects / effect_range_one_side;

      for (auto &ob_yaw_i : dynamic_objects_yaw) {
        double delt_yaw_abs = deltYawAbs(ob_yaw_i, head_yaw_plan_temp);

        if (delt_yaw_abs <= heading_resolution * effect_range_one_side) {
          double cost_dynamic_this = cost_one_piece * (effect_range_one_side - deltYawAbs(ob_yaw_i, head_yaw_plan_temp) / heading_resolution);
          cost_dynamic_objects -= cost_dynamic_this;
        }
      }
      cost_dynamic_objects = std::max(cost_dynamic_objects, cost_dynamic_objects_min);

      double cost_total_temp = cost_current_velocity + cost_planned_direction + cost_head_fluctuation + cost_dynamic_objects + cost_update_degree;

      /*** For visualization **/
      if (if_publish_panel_arrays) {
        cost_update_array[i] = cost_update_degree;
        cost_current_velocity_array[i] = cost_current_velocity;
        cost_planned_direction_array[i] = cost_planned_direction;
        cost_head_fluctuation_array[i] = cost_head_fluctuation;
        cost_dynamic_objects_array[i] = cost_dynamic_objects;
        cost_total_array[i] = cost_total_temp;
      }

      if (cost_total_temp < min_head_plan_cost)
      {
        min_head_plan_cost = cost_total_temp;
        yaw_to_send = head_yaw_plan_temp;
      }
    }

    last_head_yaw_plan = yaw_to_send;

    return yaw_to_send;
  }

  void getVisualizeCostArrays(std_msgs::Float64MultiArray &cost_array_msg,
                              std_msgs::Float64MultiArray &cost_current_velocity_array_msg,
                              std_msgs::Float64MultiArray &cost_planned_direction_array_msg,
                              std_msgs::Float64MultiArray &cost_head_fluctuation_array_msg,
                              std_msgs::Float64MultiArray &cost_dynamic_objects_array_msg,
                              std_msgs::Float64MultiArray &cost_total_array_msg)
  {
    cost_array_msg.data.clear();
    cost_current_velocity_array_msg.data.clear();
    cost_planned_direction_array_msg.data.clear();
    cost_head_fluctuation_array_msg.data.clear();
    cost_dynamic_objects_array_msg.data.clear();
    cost_total_array_msg.data.clear();

    for (int i = 0; i < LOOKING_PIECES_SIZE; i++) {
      cost_array_msg.data.push_back(cost_update_array[i]);
      cost_current_velocity_array_msg.data.push_back(cost_current_velocity_array[i]);
      cost_planned_direction_array_msg.data.push_back(cost_planned_direction_array[i]);
      cost_head_fluctuation_array_msg.data.push_back(cost_head_fluctuation_array[i]);
      cost_dynamic_objects_array_msg.data.push_back(cost_dynamic_objects_array[i]);
      cost_total_array_msg.data.push_back(cost_total_array[i]);
    }
  }

private:
  const double HEAD_BUFFER_HIT_INCREASE = 0.4;
  // Miss add value when velocity is 1m/s
  const double HEAD_BUFFER_MISS_DECREASE_STANDARD_V = 0.05;
  const double HEAD_BUFFER_MISS_DECREASE_MIN = 0.03;

  double heading_resolution = M_PI * 2 / (double)LOOKING_PIECES_SIZE;
  int mid_seq_num = (int)(LOOKING_PIECES_SIZE / 2);
  int valid_piece_num_one_side;

  // Parameters for the optimization
  double MAX_V_XY = 1.0;

  Eigen::VectorXd _direction_update_buffer =
      Eigen::VectorXd::Zero(LOOKING_PIECES_SIZE); // Range [0, 1]

  Eigen::Vector3d v0 = Eigen::Vector3d::Zero();

  // variables for the optimization
  double v_direction = 0.0;
  double last_head_yaw_plan = 0.0;
  std::vector<double> dynamic_objects_yaw;

  // cost arrays for visualization
  std::vector<double> cost_update_array{std::vector<double>(LOOKING_PIECES_SIZE, 0.0)};
  std::vector<double> cost_current_velocity_array{std::vector<double>(LOOKING_PIECES_SIZE, 0.0)};
  std::vector<double> cost_planned_direction_array{std::vector<double>(LOOKING_PIECES_SIZE, 0.0)};
  std::vector<double> cost_head_fluctuation_array{std::vector<double>(LOOKING_PIECES_SIZE, 0.0)};
  std::vector<double> cost_dynamic_objects_array{std::vector<double>(LOOKING_PIECES_SIZE, 0.0)};
  std::vector<double> cost_total_array{std::vector<double>(LOOKING_PIECES_SIZE, 0.0)};

  /// private functions

  double getHeadingYawFromSeq(int seq)
  {
    if (seq >= LOOKING_PIECES_SIZE || seq < 0) {
      std::cerr << "Seq for yaw buffer of the head out of range." << std::endl;
      return 0.0;
    }
    else {
      return (seq - mid_seq_num) * heading_resolution;
    }
  }

  int getHeadingSeq(double direction)
  {
    if (direction > M_PI) {
      direction -= M_PI * 2;
    }
    else if (direction < -M_PI) {
      direction += M_PI * 2;
    }

    if (direction > 0) { // To sovle the truncation problem
      direction += heading_resolution / 2.0;
    }
    else {
      direction -= heading_resolution / 2.0;
    }
    /// Caution!!!!! the rotating range of the camera can only be within (-PI-heading_resolution, PI+heading_resolution), chg
    int heading_seq = (int)(direction / heading_resolution) + mid_seq_num;
    if (heading_seq == LOOKING_PIECES_SIZE)
      heading_seq = 0;
    else if (heading_seq < 0)
      heading_seq = LOOKING_PIECES_SIZE - 1;

    return heading_seq;
  }

  void correctPieceSeq(int &seq)
  {
    if (seq >= LOOKING_PIECES_SIZE) { // to form a ring
      seq = seq - LOOKING_PIECES_SIZE;
    }
    else if (seq < 0) {
      seq = seq + LOOKING_PIECES_SIZE;
    }
  }

  void addHitOnePiece(int seq)
  {
    _direction_update_buffer(seq) + HEAD_BUFFER_HIT_INCREASE <= 1.0
        ? _direction_update_buffer(seq) += HEAD_BUFFER_HIT_INCREASE
        : _direction_update_buffer(seq) = 1.0;
  }

  void addMissOnePiece(int seq)
  {
    double delt_miss = HEAD_BUFFER_MISS_DECREASE_MIN;
    double delt_miss_by_velocity =
        HEAD_BUFFER_MISS_DECREASE_STANDARD_V * std::max(fabs(v0(0)), fabs(v0(1)));

    if (delt_miss_by_velocity > delt_miss) {
      delt_miss = delt_miss_by_velocity;
    }

    _direction_update_buffer(seq) - delt_miss >= 0.0
        ? _direction_update_buffer(seq) -= delt_miss
        : _direction_update_buffer(seq) = 0.0;
  }

  /** Delt yaw calculation. Avoid +=Pi problem **/
  double deltYawAbs(double yaw1, double yaw2)
  {
    double delt_yaw = fabs(yaw1 - yaw2);
    if (delt_yaw > M_PI) {
      delt_yaw = M_PI * 2 - delt_yaw;
    }
    return delt_yaw;
  }
};

#endif